import pandas as pd 
import os
import numpy as np 
import warnings
warnings.filterwarnings('ignore')

# 策略逻辑 截面开仓
# 阈值以上计算排名，买入n只，开仓至金额不够

# 买入仓位表， 当日开盘价买入
# 卖出卖出表， 当日收盘价卖出
# 持仓表

# 到期平仓判断
# 开仓判断
# 止盈止损

#阈值加rank开仓方式，附带止盈止损
class Pos_make_rank_thred_step():   
    
    def __init__(self, args, open_root=None, close_root=None):
        self.args = args
        self.money_init = args.money_init
        self.money = args.money_init
        self.money = args.money_init
        self.thred = args.open_thred
        self.vol_num = 20
        self.predict = self.get_predict()   
        self.open, self.close = self.get_price(open_root, close_root)
        self.hold_pos = self.pos_init()
        self.buy_pos = self.pos_init()
        self.sell_pos = self.pos_init()
        self.cost_df= self.pos_init()
        self.money_df = pd.DataFrame(index=self.predict.index, columns=['money'])
        self.stop_ret = args.stop_ret
        self.stop_loss = args.stop_loss
        self.cost_ratio = args.cost_ratio
        self.hold_num = args.hold_num
        self.pos_root = self.pos_root_make()

    
    
    def pos_root_make(self):
        result_root = os.path.join(self.args.result_root, self.__class__.__name__, f'loss{self.stop_loss}_ret{self.stop_ret}_thred{self.thred}_holdnum{self.hold_num}')
        if not os.path.exists(result_root):
            os.makedirs(result_root)
        return result_root
    
    def log_write(self, content):
        with open(os.path.join(self.pos_root, 'log.txt'),"a") as file:   #只需要将之前的”w"改为“a"即可，代表追加内容
            file.write(content)

    def get_price(self, open_root, close_root):  
        # sourcery skip: avoid-builtin-shadow
        open_root = r'data/stock_data/consentrate_daily_price/open.pkl.gzip'
        close_root = r'data/stock_data/consentrate_daily_price/close.pkl.gzip'
        open = pd.read_pickle(open_root)
        close = pd.read_pickle(close_root)
        open = open.loc[self.predict.index, self.predict.columns]
        close = close.loc[self.predict.index, self.predict.columns]
        
        open = open.fillna(method='ffill')
        close = close.fillna(method='ffill')
        return open, close

    def get_predict(self):
        df = pd.read_pickle(os.path.join(self.args.result_root, 'test_pred.pkl.gzip'))
        df.index = pd.to_datetime(df.index) 
        df.dropna(how='any', axis=0, inplace=True)   
        df = df.shift(1) # 预测值滞后一天
        df = df[(df.index>=self.args.backtest_time['beg'])&(df.index<=self.args.backtest_time['end'])]
        return df
    
    def pos_init(self):
        contract = self.predict.columns
        zero_ary = np.zeros((len(self.predict),len(contract)))
        pos_df = pd.DataFrame(zero_ary, index= self.predict.index, columns=contract)
        pos_df = pos_df.astype(int)
        return pos_df
    
    def pos_make(self):         
        for idx, idv in enumerate(self.predict.index):
            # 到期平仓判断
            # self.time_stop_pos(idx, idv)

            # 开仓判断
            print(idv)
            self.log_write(f'\n{idv}')
            self.pos_add(idx, idv)
            # self.pos_sub(idx, idv)
            self.stop_ratio(idx, idv)
            self.money_df.loc[idv] = self.money
            self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        self.pos_result_save()
    
    def pos_result_save(self):
        self.buy_pos.to_csv(f'{self.pos_root}/buy_pos.csv')
        self.sell_pos.to_csv(f'{self.pos_root}/sell_pos.csv')
        self.hold_pos.to_csv(f'{self.pos_root}/hold_pos.csv')
        self.money_df.to_csv(f'{self.pos_root}/money.csv')

    def time_stop_pos(self, idx, idv):
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        for code in self.columns:
            hold_pos = self.hold_pos.loc[idv, code].values
            if  hold_pos > 0:
                self.sell_pos.loc[idv, code] = hold_pos

    def pos_add(self, idx, idv):
        pred = self.predict.loc[idv]
        open_idv = (self.open.loc[idv]>=5)
        open_idv = open_idv[open_idv]
        pred = pred[open_idv.index]
        pred = pred.sort_values(ascending=False)
        pred = pred[:self.hold_num]
        buy_codes = pred.index
        buy_dict = {code: 0 for code in buy_codes}
        price_list = self.open.loc[idv,buy_codes].values*100
        cost_lst = np.maximum(price_list*self.cost_ratio, 2)
        while (self.money > cost_lst+price_list).any():
            for code in buy_codes:#debug 检查此处             
                price_amount = self.open.loc[idv, code] * 100
                cost = np.maximum(price_amount*self.cost_ratio, 1)
                if self.money >= (price_amount+cost): #资金满足开仓要求
                    buy_dict[code] += 1
                    price = self.open.loc[idv, code]
                    self.money -= (price * 100) 
                    self.cost_df.loc[idv, code] += cost  
                    self.log_write(f'\n{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')
                    print(f'{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')
        for code, num in buy_dict.items():
            self.buy_pos.loc[idv, code] = buy_dict[code]


    def pos_sub(self, idx, idv):
        if idx <=0:
            return None
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        pred = self.predict.loc[idv]
        pred = pred.sort_values(ascending=True)
        pred = pred[:int(self.thred*len(pred))]
        idv_pre = self.hold_pos.index[idx-1]
        for code in pred.index:                  
            hold_num = self.hold_pos.loc[idv_pre, code]
            sell_num = self.sell_pos.loc[idv, code]
            if (hold_num > 0) and (sell_num==0): #持仓大于0
                self.sell_pos.loc[idv, code] = hold_num
                price = self.open.loc[idv, code]
                self.money += (price * 100 *hold_num)
                self.log_write(f'\n{idv}：{code}平仓, 剩余金额：{self.money:.2f}')
                print(f'{idv}：{code}平仓, 剩余金额：{self.money:.2f}')
                cost = np.maximum(price * 100 *hold_num*self.cost_ratio, 1)
                self.cost_df.loc[idv, code] += cost
                            
    def stop_ratio(self, idx, idv): #计算平均成本
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        idv_pre = self.hold_pos.index[idx-1]
        for code in self.buy_pos.columns:
            if idx == 0:
                break
            id_init = idx - 1
            pos_i = self.buy_pos.loc[self.buy_pos.index[id_init],code]
            while (pos_i == 0) and (id_init > 0):               
                id_init = id_init - 1
                pos_i = self.buy_pos.loc[self.buy_pos.index[id_init],code]
            if pos_i == 0:
                continue

            id_initv = self.predict.index[id_init]
            price_init = self.open.loc[id_initv, code]
            price_now = self.open.loc[idv, code]
            ratio = 0
            if pos_i >= 1:
                ratio = price_now/price_init -1
            else:
                self.log_write('\nerror: 仓位为0，检查逻辑') 
                print('error: 仓位为0，检查逻辑')
            hold_open = self.open[[code]].iloc[id_init:idx]
            # #止损平仓
            # 次日止损
            if len(hold_open)<=1:
                if (ratio <= self.stop_loss):          
                    hold_num = self.hold_pos.loc[idv_pre, code]
                    sell_num = self.sell_pos.loc[idv, code]
                    if (hold_num > 0) and (sell_num==0): #持仓大于0
                        self.sell_pos.loc[idv, code] = hold_num
                        price = self.open.loc[idv, code]
                        self.money += (price * 100 *hold_num)
                        self.log_write(f'\n{idv}：{code}触发次日止损平仓, 开仓日期{id_initv}, 收益率：{ratio:.2f}')
                        print(f'{idv}：{code}触发次日止损平仓, 开仓日期{id_initv}, 收益率：{ratio:.2f}')
                        cost = np.maximum(price * 100 *hold_num*self.cost_ratio, 1)
                        self.cost_df.loc[idv, code] += cost
            else:
                price_max = np.max(hold_open.values)
                #持有出现上涨
                if price_max>price_init:
                    ratio_back = price_now/price_max -1
                    ratio_max = price_max/price_init -1
                    back_limit = np.maximum(-0.2*ratio_max, -0.2)
                    if (ratio <= 0) or (ratio_back<back_limit):          
                        hold_num = self.hold_pos.loc[idv_pre, code]
                        sell_num = self.sell_pos.loc[idv, code]
                        if (hold_num > 0) and (sell_num==0): #持仓大于0
                                self.sell_pos.loc[idv, code] = hold_num
                                price = self.open.loc[idv, code]
                                self.money += (price * 100 *hold_num)
                                self.log_write(f'\n{idv}：{code}触发回踩止盈平仓, 开仓日期{id_initv}, 收益率：{ratio:.2f}, 回踩:{ratio_back:.2f}')
                                print(f'{idv}：{code}触发回踩止盈平仓, 开仓日期{id_initv}, 收益率：{ratio:.2f}, 回踩:{ratio_back:.2f}')
                                cost = np.maximum(price * 100 *hold_num*self.cost_ratio, 1)
                                self.cost_df.loc[idv, code] += cost
                    # 持有仅下跌
                if (price_max <= price_init) and (ratio <= self.stop_loss):
                    hold_num = self.hold_pos.loc[idv_pre, code]
                    sell_num = self.sell_pos.loc[idv, code]
                    if (hold_num > 0) and (sell_num==0): #持仓大于0
                        self.sell_pos.loc[idv, code] = hold_num
                        price = self.open.loc[idv, code]
                        self.money += (price * 100 *hold_num)
                        self.log_write(f'\n{idv}：{code}触发下跌止损平仓, 开仓日期{id_initv}, 收益率：{ratio:.2f}')
                        print(f'{idv}：{code}触发下跌止损平仓, 开仓日期{id_initv}, 收益率：{ratio:.2f}')
                        cost = np.maximum(price * 100 *hold_num*self.cost_ratio, 1)
                        self.cost_df.loc[idv, code] += cost
                   
        

class Pos_make_rank_thred_month(Pos_make_rank_thred_step):

    def __init__(self, args, open_root=None, close_root=None):
        super().__init__(args, open_root, close_root)

    def get_predict(self):    
        pred_list = []
        file_list = sorted(os.listdir(self.args.result_root))
        for i in file_list:
            if i[-2:] == '01':
                df_i = pd.read_pickle(os.path.join(self.args.result_root, i, 'test_pred.pkl.gzip'))
                pred_list.append(df_i)
        df = pd.concat(pred_list, axis=0)   
        df.index = pd.to_datetime(df.index) 
        df.dropna(how='any', axis=0, inplace=True)   
        df = df.shift(1) # 预测值滞后一天
        df = df[(df.index>=self.args.backtest_time['beg'])&(df.index<=self.args.backtest_time['end'])]
        return df

    # def pos_root_make(self):
    #     result_root = os.path.join(self.args.result_root, Pos_make_rank_thred_month.__name__)
    #     return result_root


class Pos_make_series_thred_month(Pos_make_rank_thred_month):  

    def __init__(self, args, open_root=None, close_root=None):
        super().__init__(args, open_root, close_root) 
        self.predict_vol = self.predict.rolling(self.vol_num, min_periods=1).std()   
        self.vol_num = 20
        self.thred = args.open_thred

    def pos_add(self, idx, idv):
        pred = self.predict.loc[idv] / (self.predict_vol.loc[idv]+0.0000001)
        open_idv = (self.open.loc[idv]>=5)
        open_idv = open_idv[open_idv]
        pred = pred[open_idv.index]
        pred = pred.sort_values(ascending=False)

        pred = pred[:self.hold_num]
        buy_codes = pred.index
        buy_dict = {code: 0 for code in buy_codes}
        price_list = self.open.loc[idv,buy_codes].values*100
        cost_lst = np.maximum(price_list*self.cost_ratio, 2)
        while (self.money > cost_lst+price_list).any():
            for code in buy_codes:#debug 检查此处
                # if pred[code] >= self.thred: #满足开仓信号要求                       
                price_amount = self.open.loc[idv, code] * 100
                cost = np.maximum(price_amount*self.cost_ratio, 1)
                if self.money >= (price_amount+cost): #资金满足开仓要求
                    buy_dict[code] += 1
                    price = self.open.loc[idv, code]
                    self.money -= (price * 100) 
                    self.cost_df.loc[idv, code] += cost  
                    self.log_write(f'\n{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')
                    print(f'{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')
        for code, num in buy_dict.items():
            self.buy_pos.loc[idv, code] = buy_dict[code]
                
    # def pos_root_make(self):
    #     result_root = os.path.join(self.args.result_root, Pos_make_series_thred_month.__name__)
    #     return result_root


class Pos_make_series_thred_step(Pos_make_series_thred_month):
    
    def __init__(self, args, open_root=None, close_root=None):
        super().__init__(args, open_root, close_root)

    def get_predict(self):
        df = pd.read_pickle(os.path.join(self.args.result_root, 'test_pred.pkl.gzip'))
        df.index = pd.to_datetime(df.index) 
        df.dropna(how='any', axis=0, inplace=True)   
        df = df.shift(1) # 预测值滞后一天
        df = df[(df.index>=self.args.backtest_time['beg'])&(df.index<=self.args.backtest_time['end'])]
        return df

    # def pos_root_make(self):
    #     result_root = os.path.join(self.args.result_root, Pos_make_series_thred_step.__name__)
    #     return result_root